# 导包
from torch.autograd import Variable
import numpy as np
import torch
import matplotlib.pyplot as plt
import torch.nn as nn
import math
import copy
import torch.nn.functional as F


# 输入部分
# 1 词嵌入层
class Embedding(nn.Module):
    def __init__(self, vocab_size, Embedding_dim):
        super(Embedding, self).__init__()
        self.vocab_size = vocab_size
        self.Embedding_dim = Embedding_dim
        self.emd = nn.Embedding(self.vocab_size, self.Embedding_dim)

    def forward(self, x):
        out = self.emd(x) * math.sqrt(self.Embedding_dim)
        return out


# 2 位置编码
class PositionalEncoding(nn.Module):
    def __init__(self, embedding_dim, drop_out=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.embedding_dim = embedding_dim
        self.drop = nn.Dropout(drop_out)
        self.max_len = max_len
        pe = torch.zeros(self.max_len, self.embedding_dim)
        position = torch.arange(0, self.max_len).unsqueeze(1)
        # div_term = torch.exp(torch.arange(0, self.embedding_dim, 2) * (-math.log(10000) / self.embedding_dim))
        # pe[:, 1::2] = torch.sin(position * div_term)
        # pe[:, 0::2] = torch.cos(position * div_term)
        _i = torch.arange(0, embedding_dim, step=2)
        pe[:, 0::2] = torch.sin(position / (10000 ** (_i / embedding_dim)))
        pe[:, 1::2] = torch.cos(position / (10000 ** (_i / embedding_dim)))

        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=True)
        return self.drop(x)


# 输入部分
# 1 构造 掩码辅助函数
def mask(size):
    shape = (1, size, size)
    s_mask = np.triu(np.ones(shape), k=-1).astype('uint8')
    x_mask = torch.from_numpy(1 - s_mask)
    return x_mask


# 2 attention
def attention(q, k, v, mask=None, dropout=None):
    embedding_dim = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(embedding_dim)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, 6)
    if dropout is not None:
        scores = dropout(scores)
    attn_S = torch.softmax(scores, dim=-1)
    out = torch.matmul(scores, v)
    return out, attn_S


# 3 定义克隆函数
def clones(module, n):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(n)])


# 4 MultiHeadAttention 多头注意机制
class MultiHeadAttention(nn.Module):
    def __init__(self, head, embedding_dim, drop_out=0.1):
        super(MultiHeadAttention, self).__init__()
        assert embedding_dim % head == 0
        self.d_k = embedding_dim // head
        self.head = head
        self.linears = clones(nn.Linear(embedding_dim, embedding_dim), 4)
        self.embedding_dim = embedding_dim
        self.attn = None
        self.drop_out = nn.Dropout(drop_out)

    def forward(self, q, k, v, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(0)
        batch_size = q.size(0)
        q, k, v = [model(x).view(batch_size, -1, self.head, self.d_k).transpose(1, 2) for model, x in
                   zip(self.linears, (q, k, v))]
        x, self.attn = attention(q, k, v, mask=mask, dropout=self.drop_out)
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.head * self.d_k)
        return self.linears[-1](x)


# 5 前馈全连接层
class PositionWiseForword(nn.Module):
    def __init__(self, embedding_dim, d_ff, drop_out=0.1):
        super(PositionWiseForword, self).__init__()
        self.embedding_dim = embedding_dim
        self.d_ff = d_ff
        self.linear1 = nn.Linear(embedding_dim, d_ff)
        self.linear2 = nn.Linear(d_ff, embedding_dim)
        self.relu = nn.ReLU()
        self.d = nn.Dropout(drop_out)

    def forward(self, x):
        return self.linear2(self.d(self.relu(self.linear1(x))))


# 6 规范层
class LayerNorm(nn.Module):
    def __init__(self, feature, eps=1e-8):
        super(LayerNorm, self).__init__()
        self.a = nn.Parameter(torch.ones(feature))
        self.b = nn.Parameter(torch.zeros(feature))
        self.f = feature
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        out = self.a * (x - mean) / (std + self.eps) + self.b
        return out


# 7 子层连接结构 多头或前馈 加残差 加规范
class SublayerConnection(nn.Module):
    def __init__(self, size, drop_out):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.drop = nn.Dropout(drop_out)

    def forward(self, x, sub):
        return self.drop(sub(self.norm(x))) + x


# 8 编码器层
class EncoderLyaer(nn.Module):
    def __init__(self, size, self_attn, self_forward, d=0.1):
        super(EncoderLyaer, self).__init__()
        self.size = size
        self.attn = self_attn
        self.fw = self_forward
        self.subs = clones(SublayerConnection(self.size, d), 2)

    def forward(self, x, m=None):
        x = self.subs[0](x, lambda x: self.attn(x, x, x, m))
        x = self.subs[1](x, self.fw)
        return x


# 9 编码器
class Encoder(nn.Module):
    def __init__(self, layer, n):
        super(Encoder, self).__init__()
        self.layers = clones(layer, n)
        self.norm = LayerNorm(layer.size)

    def forward(self, x, mask=None):
        for Layer in self.layers:
            x = Layer(x, mask)
        return self.norm(x)


# 10 解码器层
class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, self_no_attn, forward, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.attn = self_attn
        self.no_attn = self_no_attn
        self.fw = forward
        self.layers = clones(SublayerConnection(self.size, dropout), 3)

    def forward(self, x, memory, attn_mask=None, no_attn_mask=None):
        m = memory
        x = self.layers[0](x, lambda x: self.attn(x, x, x, attn_mask))
        x = self.layers[1](x, lambda x: self.no_attn(x, m, m, no_attn_mask))
        x = self.layers[2](x, self.fw)
        return x


# 11 解码器
class Decoder(nn.Module):
    def __init__(self, lyaer, n):
        super(Decoder, self).__init__()
        self.layers = clones(lyaer, 6)
        self.norm = LayerNorm(lyaer.size)

    def forward(self, x, m, attn_m=None, no_attn_m=None):
        for layer in self.layers:
            x = layer(x, m, attn_m, no_attn_m)
        return self.norm(x)


# 输出层
class Genertor(nn.Module):
    def __init__(self, embedding_dim, out_dim):
        super(Genertor, self).__init__()
        self.linear = nn.Linear(embedding_dim, out_dim)

    def forward(self, x):
        return F.log_softmax(self.linear(x), dim=-1)


# Transformer
class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, yuan_embed, mu_embed, gen):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.yuan_embed = yuan_embed
        self.mu_embed = mu_embed
        self.gen = gen

    def forward(self, yuan, mu, yuan_mask, mu_mask):
        return self.gen(self.decode(self.encode(yuan, yuan_mask), yuan_mask, mu, mu_mask))

    def encode(self, yuan, yuan_mask):
        emd_yaun = self.yuan_embed(yuan)
        return self.encoder(emd_yaun, yuan_mask)

    def decode(self, m, yuan_mask, mu, mu_mask):
        embed_mu = self.mu_embed(mu)
        return self.decoder(embed_mu, m, yuan_mask, mu_mask)


def make_model(yuan_dim, mu_dim, embedding_dim=512, n=6, d_ff=2048, head=8, dropout=0.1):
    c = copy.deepcopy
    mha = MultiHeadAttention(head, embedding_dim)
    fw = PositionWiseForword(embedding_dim, d_ff)
    pe = PositionalEncoding(embedding_dim)
    model = EncoderDecoder(
        Encoder(EncoderLyaer(embedding_dim, c(mha), c(fw)), n),
        Decoder(DecoderLayer(embedding_dim, c(mha), c(mha), c(fw)), n),
        nn.Sequential(Embedding(yuan_dim, embedding_dim), c(pe)),
        nn.Sequential(Embedding(mu_dim, embedding_dim), c(pe)),
        Genertor(embedding_dim, mu_dim)
    )
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model


def dm_test_model():
    yuan_dim = 200
    mu_dim = 1000
    my_Transformer = make_model(yuan_dim, mu_dim)
    print(my_Transformer)
    # 假设源数据与目标数据相同, 实际中并不相同模拟
    source = target = Variable(torch.LongTensor([[1, 2, 3, 8], [3, 4, 1, 8]]))

    # 假设src_mask与tgt_mask相同，实际中并不相同
    source_mask = target_mask = mask(4)  #
    mydata = my_Transformer(source, target, source_mask, target_mask)
    print('mydata.shape--->', mydata.shape)
    print('mydata--->', mydata)


if __name__ == '__main__':
    dm_test_model()
